SageMaker StudioからSageMaker SDK経由でHugging Faceモデルを呼び出してファインチューニングを実行する

データアナリティクス事業本部 インテグレーション部 機械学習チームの貞松です。

本記事は「クラスメソッド 機械学習チーム アドベントカレンダー 2022」の22日目です。
昨日(21日目)は Shirota による 「コサイン類似度」で文書がどれだけ似ているかを調べてみた でした。

本記事では、SageMaker StudioからSageMaker SDK経由でHugging Faceモデルを呼び出してファインチューニングを実行する(正確にはHugging Face用のDeepLearning Container上でジョブを実行する)手順について解説します。

Hugging Faceとは

Hugging Faceは、機械学習を使用してアプリケーションを構築するためのツールを開発する企業です。
自然言語処理アプリケーション用に構築されたTransformersライブラリとHugging Face Hubという、ユーザーが機械学習モデルとデータセットを共有できるプラットフォームが特に有名です。

Transformersライブラリ

Transformersライブラリは、テキストや画像、音声用のTransformerモデルのオープンソース実装を含むPythonパッケージです。
PyTorch、TensorFlow、およびJAXと互換性があり、BERTやGPT等の有名なモデルの実装が含まれています。

Hugging Face Hub

Hugging Face Hubは、ユーザーがデータセットや事前トレーニング済みのモデル、機械学習プロジェクトのデモを共有できるプラットフォームです。
プロジェクトのディスカッションやプルリクエストなど、コード共有とコラボレーションに必要なGitHubライクな機能が含まれています。
また、ユーザーがGradioまたはStreamlitを使用して、機械学習アプリのWebベースのデモを構築できるホストサービスであるHugging Face Spacesを有しています。

Hugging Face on SageMakerについて

SageMaker上で実行されるHugging Faceは、SageMaker SDK経由でHugging Faceモデルに対するファインチューニングや推論の実行、デプロイを実行する為のジョブを操作します。 実際にジョブが実行される環境として、Hugging Faceモデル用のAWS Deep Learning Containersを利用します。

AWS公式のドキュメントは以下をご参照ください。

SageMakerのAPIリファレンス内のHugging Faceに関するAPIについては以下をご参照ください。

Hugging Face公式のHugging Face on SageMakerに関する解説、チュートリアル等については以下をご参照ください。

SageMaker Studio上でHugging Faceモデルのファインチューニングを実行する

Hugging Face公式がGitHub上に用意しているSageMaker用のサンプルノートブックを使用して、Hugging Faceモデルのファインチューニングを実行します。

対象とするモデルの具体的な内容としては、imdbデータセット(大規模な映画レビュー データセット)を用いた、テキストの2値分類モデル(POSITIVE or NEGATIVE)となっています。

↓imdbデータセットについて

↓使用したサンプルノートブック

SageMaker Studioを起動して、GitタブからサンプルノートブックのGitHubリポジトリをクローンします。

クローンが完了したら notebooks/sagemaker/01_getting_started_pytorch フォルダから sagemaker-notebook.ipynb を開きます。

ノートブックの実行環境として、imageに Pytorch 1.6 Python 3.6 CPU Optimized を選択します。

以降、ノートブックの中身について順番に解説していきます。

ライブラリインストール

まずは必要なライブラリのインストールです。
SageMaker SDKとHugging Faceのtransformers、datasetsライブラリをインストールしておきます。

!pip install "sagemaker>=2.48.0" "transformers==4.12.3" "datasets[s3]==1.18.3" --upgrade

環境設定や権限設定など

SageMaker SDKのHugging Face用モジュールをインポートします。
また、SageMakerのセッションを作成して、必要な権限を持つ実行ロールを割り当てます。

import sagemaker.huggingface
import sagemaker

sess = sagemaker.Session()
# sagemaker session bucket -> used for uploading data, models and logs
# sagemaker will automatically create this bucket if it not exists
sagemaker_session_bucket=None
if sagemaker_session_bucket is None and sess is not None:
    # set to default bucket if a bucket name is not given
    sagemaker_session_bucket = sess.default_bucket()

role = sagemaker.get_execution_role()
sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)

print(f"sagemaker role arn: {role}")
print(f"sagemaker bucket: {sess.default_bucket()}")
print(f"sagemaker session region: {sess.boto_region_name}")

データロードと前処理

imdbデータセットをロードして、トークン化を施します。

from datasets import load_dataset
from transformers import AutoTokenizer

# tokenizer used in preprocessing
tokenizer_name = 'distilbert-base-uncased'

# dataset used
dataset_name = 'imdb'

# s3 key prefix for the data
s3_prefix = 'samples/datasets/imdb'

# load dataset
dataset = load_dataset(dataset_name)

# download tokenizer
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)

# tokenizer helper function
def tokenize(batch):
    return tokenizer(batch['text'], padding='max_length', truncation=True)

# load dataset
train_dataset, test_dataset = load_dataset('imdb', split=['train', 'test'])
test_dataset = test_dataset.shuffle().select(range(10000)) # smaller the size for test dataset to 10k 


# tokenize dataset
train_dataset = train_dataset.map(tokenize, batched=True)
test_dataset = test_dataset.map(tokenize, batched=True)

# set format for pytorch
train_dataset =  train_dataset.rename_column("label", "labels")
train_dataset.set_format('torch', columns=['input_ids', 'attention_mask', 'labels'])
test_dataset = test_dataset.rename_column("label", "labels")
test_dataset.set_format('torch', columns=['input_ids', 'attention_mask', 'labels'])

処理済みのデータセットについて、学習用、検証用それぞれをS3バケットに配置します。

import botocore
from datasets.filesystems import S3FileSystem

s3 = S3FileSystem()  

# save train_dataset to s3
training_input_path = f's3://{sess.default_bucket()}/{s3_prefix}/train'
train_dataset.save_to_disk(training_input_path,fs=s3)

# save test_dataset to s3
test_input_path = f's3://{sess.default_bucket()}/{s3_prefix}/test'
test_dataset.save_to_disk(test_input_path,fs=s3)

実行スクリプトの確認

ノートブックと同じフォルダに同梱されている実行スクリプト(ファインチューニングのトレーニングジョブで実行される処理の本体)の中身を確認します。
トレーニングに関する設定とデータの入出力先に関する設定を渡すことで、それらを使用したモデル学習の処理と結果の保存を実行してくれるようです。

!pygmentize ./scripts/train.py

from transformers import AutoModelForSequenceClassification, Trainer, TrainingArguments, AutoTokenizer
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
from datasets import load_from_disk
import random
import logging
import sys
import argparse
import os
import torch

if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    # hyperparameters sent by the client are passed as command-line arguments to the script.
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--train_batch_size", type=int, default=32)
    parser.add_argument("--eval_batch_size", type=int, default=64)
    parser.add_argument("--warmup_steps", type=int, default=500)
    parser.add_argument("--model_name", type=str)
    parser.add_argument("--learning_rate", type=str, default=5e-5)

    # Data, model, and output directories
    parser.add_argument("--output_data_dir", type=str, default=os.environ["SM_OUTPUT_DATA_DIR"])
    parser.add_argument("--model_dir", type=str, default=os.environ["SM_MODEL_DIR"])
    parser.add_argument("--n_gpus", type=str, default=os.environ["SM_NUM_GPUS"])
    parser.add_argument("--training_dir", type=str, default=os.environ["SM_CHANNEL_TRAIN"])
    parser.add_argument("--test_dir", type=str, default=os.environ["SM_CHANNEL_TEST"])

    args, _ = parser.parse_known_args()

    # Set up logging
    logger = logging.getLogger(__name__)

    logging.basicConfig(
        level=logging.getLevelName("INFO"),
        handlers=[logging.StreamHandler(sys.stdout)],
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )

    # load datasets
    train_dataset = load_from_disk(args.training_dir)
    test_dataset = load_from_disk(args.test_dir)

    logger.info(f" loaded train_dataset length is: {len(train_dataset)}")
    logger.info(f" loaded test_dataset length is: {len(test_dataset)}")

    # compute metrics function for binary classification
    def compute_metrics(pred):
        labels = pred.label_ids
        preds = pred.predictions.argmax(-1)
        precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average="binary")
        acc = accuracy_score(labels, preds)
        return {"accuracy": acc, "f1": f1, "precision": precision, "recall": recall}

    # download model from model hub
    model = AutoModelForSequenceClassification.from_pretrained(args.model_name)
    tokenizer = AutoTokenizer.from_pretrained(args.model_name)

    # define training args
    training_args = TrainingArguments(
        output_dir=args.model_dir,
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.train_batch_size,
        per_device_eval_batch_size=args.eval_batch_size,
        warmup_steps=args.warmup_steps,
        evaluation_strategy="epoch",
        logging_dir=f"{args.output_data_dir}/logs",
        learning_rate=float(args.learning_rate),
    )

    # create Trainer instance
    trainer = Trainer(
        model=model,
        args=training_args,
        compute_metrics=compute_metrics,
        train_dataset=train_dataset,
        eval_dataset=test_dataset,
        tokenizer=tokenizer,
    )

    # train model
    trainer.train()

    # evaluate model
    eval_result = trainer.evaluate(eval_dataset=test_dataset)

    # writes eval result to file which can be accessed later in s3 ouput
    with open(os.path.join(args.output_data_dir, "eval_results.txt"), "w") as writer:
        print(f"***** Eval results *****")
        for key, value in sorted(eval_result.items()):
            writer.write(f"{key} = {value}\n")

    # Saves the model to s3
    trainer.save_model(args.model_dir)

Estimatorの設定とファインチューニングの実行

学習ジョブで使用するハイパーパラメータを設定します。

from sagemaker.huggingface import HuggingFace

# hyperparameters, which are passed into the training job
hyperparameters={'epochs': 1,
                 'train_batch_size': 32,
                 'model_name':'distilbert-base-uncased'
                 }

Hugging Face Estimatorを作成してトレーニングジョブを実行します。
実行インスタンスの設定について、オリジナルのコードでは ml.p3.2xlarge を設定しているのですが、正直少しサイズダウンして ml.g4dn.2xlarge に変更して実行しても1時間も掛からずに実行が完了するので、ここは変更しても良いと思います(時間単価は約1/3になります)

huggingface_estimator = HuggingFace(entry_point='train.py',
                            source_dir='./scripts',
                            instance_type='ml.g4dn.2xlarge', # 元はml.p3.2xlarge
                            instance_count=1,
                            role=role,
                            transformers_version='4.12',
                            pytorch_version='1.9',
                            py_version='py38',
                            hyperparameters = hyperparameters)

設定したEstimatorでトレーニングジョブの実行を開始します。
※学習実行時のログはハチャメチャに長いので、ここでは割愛します。

huggingface_estimator.fit({'train': training_input_path, 'test': test_input_path})

エンドポイントデプロイと推論の実行

トレーニングジョブの実行が完了したら、エンドポイントデプロイを実行します。

predictor = huggingface_estimator.deploy(1,"ml.g4dn.xlarge")

作成されたエンドポイントにアクセスして、推論結果を取得します。

sentiment_input= {"inputs":"I love using the new Inference DLC."}

predictor.predict(sentiment_input)

ラベルとスコアが返ってきます。
LABEL_1はPOSITIVEに対応するラベルですね。どうやら正常に動作しているようです。

[{'label': 'LABEL_1', 'score': 0.8758626580238342}]

確認が終わったらエンドポイントは削除します。

predictor.delete_endpoint()

Extras. 以前実行されたトレーニングジョブのログ確認や学習済みモデルのロードを実行する

追加の確認として、実行済みのトレーニングジョブのログやその際の学習済みモデルをロードする便利機能を確認します。

まずは前述で実行されたトレーニングジョブの情報を直接取り出して中身を確認しておきます。
後々使用する為、トレーニングジョブ名は変数に格納しておきます。

# container image used for training job
print(f"container image used for training job: \n{huggingface_estimator.image_uri}\n")

# s3 uri where the trained model is located
print(f"s3 uri where the trained model is located: \n{huggingface_estimator.model_data}\n")

# latest training job name for this estimator
latest_training_job_name = huggingface_estimator.latest_training_job.name
print(f"latest training job name for this estimator: \n{latest_training_job_name}\n")

sagemaker_session.logs_for_job にトレーニングジョブ名を渡すことで、ジョブ実行時のログを一式取得することができます。
※学習実行時のログはハチャメチャに長いので、ここでは割愛します。
オリジナルのコードでは、トレーニングジョブ名として huggingface_estimator.latest_training_job.name を渡していますが、前述で変数に格納済みの為、これを利用します。

huggingface_estimator.sagemaker_session.logs_for_job(latest_training_job_name)

オリジナルのコードでは old_training_job_name が空になっていますが、ここも上記で格納しておいた変数を渡すことで正しくモデルをロードすることができます。
これにより、ロードしたモデルをコード上で操作したり、モデルの実体がどこのS3パスに格納されているか確認したりできます。

from sagemaker.estimator import Estimator

# job which is going to be attached to the estimator
old_training_job_name=latest_training_job_name

# attach old training job
huggingface_estimator_loaded = Estimator.attach(old_training_job_name)

# get model output s3 from training job
huggingface_estimator_loaded.model_data

まとめ

Hugging Face on SageMakerに関する概要と、SageMaker Studio上で実際にHugging Faceモデルのファインチューニングを実行する手順について解説しました。
SageMakerは、柔軟にコンピューティング環境を切り替えたり、スケールすることができる環境を即時用意することができるので、既存の学習済みモデルを活用した実験や実装と相性が良さそうです。
Hugging Faceと言えばTransformersライブラリということで、テキストや画像、音声に関する推論モデルをサクッと試してみたい方のお役に立てば幸いです。